package com.liru.neo4j_demo.controller;

import com.liru.neo4j_demo.util.FileSplitter;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.store.embedding.*;
import jakarta.annotation.Resource;
import org.apache.commons.csv.CSVFormat;
import org.apache.commons.csv.CSVParser;
import org.apache.commons.csv.CSVRecord;
import org.neo4j.driver.*;
import org.neo4j.driver.Record;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;

import java.io.*;
import java.nio.charset.StandardCharsets;
import java.util.*;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * AController 控制器类，用于处理与嵌入和图数据库相关的请求
 */
@RestController
@RequestMapping("/a")
public class AController {
    // 注入嵌入存储实例，用于存储和检索文本片段的嵌入向量
    @Resource
    EmbeddingStore<TextSegment> embeddingStore;
    
    // 注入嵌入模型实例，用于将文本转换为向量表示
    @Resource
    EmbeddingModel embeddingModel;

    // 注入Neo4j数据库驱动程序，用于与图数据库交互
    private final Driver neo4jDriver = GraphDatabase.driver("neo4j://192.168.6.63:7687", AuthTokens.basic("neo4j", "Secret123"));

    // 注入嵌入存储摄入器，用于处理文档并将其转换为嵌入向量存储
    @Resource
    private EmbeddingStoreIngestor embeddingStoreIngestor;

    // 从配置文件中读取OpenAI API密钥
    @Value("${openai.api.key:}")
    private String openaiApiKey;

    @Resource
    private OpenAiChatModel openAiChatModel;

    /**
     * 处理文件上传、自定义切割并批量存入Neo4j的请求
     * @param file 上传的文件
     * @param splitType 切割类型：lines, paragraphs, characters, regex
     * @param splitParam 切割参数：行数、字符数或正则表达式
     * @return 操作结果字符串
     * @throws IOException 文件处理异常
     */
    @PostMapping("/upload/split")
    public String uploadAndSplit(
            @RequestParam("file")MultipartFile file,
            @RequestParam("splitType")String splitType,
            @RequestParam("splitParam")String splitParam) throws IOException {
        // 获取上传文件的原始文件名
        String fileName = file.getOriginalFilename();
        // 初始化文本内容变量
        String textContent = "";

        // 判断文件是否为CSV格式   //TODO 对csv文件是否需要特殊处理？
        if (fileName != null && fileName.toLowerCase().endsWith(".csv")) {
            // 如果是CSV文件，则解析CSV内容
            textContent = parseCsvFile(file);
        } else {
            // 如果是普通文本文件，则直接读取文件内容   //TODO 读取时，考虑编码格式问题，如果乱码是否有影响？
            textContent = new String(file.getBytes(), StandardCharsets.UTF_8);
        }

        // 根据切割类型和参数切割文本
        List<String> segments = splitText(textContent, splitType, splitParam);

        // 批量将片段存入Neo4j
        batchInsertSegments(segments, fileName);

        // 返回操作成功标识
        return "成功上传并切割文件：" + fileName + "，生成" + segments.size() + "个片段"; 
    }

    /**
     * 根据切割类型和参数切割文本
     * @param text 待切割的文本
     * @param splitType 切割类型
     * @param splitParam 切割参数
     * @return 切割后的文本片段列表
     */
    private List<String> splitText(String text, String splitType, String splitParam) {   //TODO 分割函数优化？
        return switch (splitType.toLowerCase()) {
            case "lines" -> {
                // 按行数切割
                int linesPerSegment = Integer.parseInt(splitParam);
                yield FileSplitter.splitByLines(text, linesPerSegment);
            }
            case "paragraphs" ->
                // 按段落切割
                    FileSplitter.splitByParagraphs(text);
            case "characters" -> {
                // 按字符数切割
                int charsPerSegment = Integer.parseInt(splitParam);
                yield FileSplitter.splitByCharacters(text, charsPerSegment);
            }
            case "regex" ->
                // 按正则表达式切割
                    FileSplitter.splitByRegex(text, splitParam);
            default ->
                // 默认按段落切割
                    FileSplitter.splitByParagraphs(text);
        };
    }

    /**
     * 批量将文本片段插入Neo4j数据库
     * @param segments 文本片段列表
     * @param fileName 文件名
     */
    private void batchInsertSegments(List<String> segments, String fileName) { //TODO 文本插入需向量化
        try (Session session = neo4jDriver.session()) {
            // 开启事务
            session.writeTransaction(tx -> {
                AtomicInteger index = new AtomicInteger(0);
                List<String> segmentIds = new ArrayList<>();
                // 创建每个片段节点
                for (String segment : segments) {
                    int currentIndex = index.getAndIncrement();
                    // 创建节点并返回ID
                    Result result = tx.run(  //TODO 构建节点，节点实体如何选择（不同标签如：文章，文本     部门，员工，技能）？如果能够自动适配？
                            "CREATE (s:Segment {content: $content, fileName: $fileName, index: $index, createdAt: $timestamp}) " +
                            "RETURN id(s) as id",
                            Values.parameters(
                                    "content", segment,
                                    "fileName", fileName,
                                    "index", currentIndex,
                                    "timestamp", System.currentTimeMillis()
                            )
                    );

                    // 获取节点ID（长整型）   //TODO 这个插入向量是什么
                    long id = result.single().get("id").asLong();
                    segmentIds.add(String.valueOf(id));

                    // 为片段创建嵌入向量
                    Embedding embedding = embeddingModel.embed(segment).content();
                    TextSegment textSegment = TextSegment.from(segment);
                    embeddingStore.add(embedding, textSegment);
                }

                // 创建片段之间的NEXT关系
                for (int i = 0; i < segmentIds.size() - 1; i++) { //TODO 节点关系考虑？（上下文，文章所属，员工索树部门）如何自动适配？
                    tx.run(
                            "MATCH (s1), (s2) WHERE id(s1) = $id1 AND id(s2) = $id2 " +
                            "CREATE (s1)-[:NEXT_SEGMENT]->(s2)",
                            Values.parameters(
                                    "id1", segmentIds.get(i),
                                    "id2", segmentIds.get(i + 1)
                            )
                    );
                }

                return null;
            });
        } catch (Exception e) {
            throw new RuntimeException("批量插入片段到Neo4j失败: " + e.getMessage(), e);
        }
    }


    /**
     * 解析CSV文件内容
     * @param file 上传的CSV文件
     * @return 解析后的文本内容
     * @throws IOException 文件处理异常
     */
    private String parseCsvFile(MultipartFile file) throws IOException {
        // 创建字符串构建器用于存储解析后的内容
        StringBuilder content = new StringBuilder();
        
        // 创建文件读取器和CSV解析器
        try (Reader reader = new InputStreamReader(file.getInputStream(), StandardCharsets.UTF_8);
             CSVParser csvParser = new CSVParser(reader, CSVFormat.DEFAULT.withFirstRecordAsHeader())) {

            // 遍历CSV文件中的每条记录
            for (CSVRecord record : csvParser) {
                // 将每条记录转换为Map格式并添加到内容中
                content.append(record.toMap().toString()).append("\n");
            }
        }
        // 返回解析后的完整内容
        return content.toString();
    }



    /**
     * 处理智能问答请求
     * 整合向量搜索、关系分析和AI响应生成
     * @param question 用户问题
     * @return AI生成的回答
     */
    @RequestMapping("/ask")
    public String ask(String question) {
        // 1. 向量搜索相关片段
        List<EmbeddingMatch<TextSegment>> matches = searchRelevantSegments(question);

        if (matches.isEmpty()) {
            return "未找到相关信息来回答您的问题。";
        }

        // 2. 从Neo4j获取相关节点和关系
        List<Map<String, Object>> relatedInfo = getRelatedInfoFromNeo4j(matches);

        // 3. 构造提示词
        String prompt = buildPrompt(question, relatedInfo);

        // 4. 发送给AI模型并获取响应
        return getAIResponse(prompt);
    }

    /**
     * 向量搜索相关片段
     * @param question 用户问题
     * @return 匹配的片段列表
     */
    private List<EmbeddingMatch<TextSegment>> searchRelevantSegments(String question) {
        Embedding queryEmbed = embeddingModel.embed(question).content();

        EmbeddingSearchRequest embeddingRequest = EmbeddingSearchRequest.builder()
                .queryEmbedding(queryEmbed)
                .minScore(0.5)
                .maxResults(5)
                .build();
        //TODO 向量搜索方法1   如何让它可以根据关系衍生搜索
        return embeddingStore.search(embeddingRequest).matches();
    }

    /**
     * 从Neo4j获取相关节点和关系
     * @param matches 向量搜索匹配结果
     * @return 相关信息列表
     */
    private List<Map<String, Object>> getRelatedInfoFromNeo4j(List<EmbeddingMatch<TextSegment>> matches) {
        List<Map<String, Object>> relatedInfo = new ArrayList<>();

        try (Session session = neo4jDriver.session()) {
            for (EmbeddingMatch<TextSegment> match : matches) {
                String segmentText = match.embedded().text();

                // 查询包含该文本的节点及其关系
                Result result = session.run(
                        "MATCH (s:Segment)-[r]->(related) WHERE s.content CONTAINS $text " +
                        "RETURN s.content AS segment, type(r) AS relation, related.content AS relatedContent",
                        Values.parameters("text", segmentText.substring(0, Math.min(50, segmentText.length())))
                );

                while (result.hasNext()) {
                    Record record = result.next();
                    Map<String, Object> info = new HashMap<>();
                    info.put("segment", record.get("segment").asString());
                    info.put("relation", record.get("relation").asString());
                    info.put("relatedContent", record.get("relatedContent").asString());
                    relatedInfo.add(info);
                }
            }
        } catch (Exception e) {
            System.err.println("从Neo4j获取相关信息失败: " + e.getMessage());
        }

        return relatedInfo;
    }

    /**
     * 构造提示词
     * @param question 用户问题
     * @param relatedInfo 相关信息
     * @return 构造好的提示词
     */
    private String buildPrompt(String question, List<Map<String, Object>> relatedInfo) {
        StringBuilder prompt = new StringBuilder();
        prompt.append("基于以下信息回答问题：\n\n");

        if (!relatedInfo.isEmpty()) {
            prompt.append("相关上下文信息：\n");
            for (Map<String, Object> info : relatedInfo) {
                prompt.append("- 片段: ").append(info.get("segment")).append("\n")
                        .append("  关系: ").append(info.get("relation")).append("\n")
                        .append("  相关内容: ").append(info.get("relatedContent")).append("\n\n");
            }
        } else {
            prompt.append("没有找到相关上下文信息。\n\n");
        }

        prompt.append("问题: ").append(question).append("\n\n");
        prompt.append("请根据提供的信息，简洁、准确地回答问题。如果信息不足，请说明。");

        return prompt.toString();
    }

    /**
     * 获取AI模型的响应
     * @param prompt 提示词
     * @return AI生成的回答
     */
    private String getAIResponse(String prompt) {
        try {
            // 发送提示词并获取响应
            return openAiChatModel.chat(prompt);
        } catch (Exception e) {
            return "AI响应生成失败: " + e.getMessage();
        }
    }
}
